import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset

class AudioCNN(nn.Module):
    def __init__(self):
        super(AudioCNN, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
        self.fc1 = nn.Linear(32 * 64 * 64, 128)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(128, 4)  # 假设有4种情感类别

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = x.view(-1, 32 * 64 * 64)  # 展平
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

class TextEmbedding(nn.Module):
    def __init__(self, embedding_dim, vocab_size, text_length):
        super(TextEmbedding, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, 64, batch_first=True)
        self.fc = nn.Linear(64, 128)

    def forward(self, x):
        x = self.embedding(x)
        x, _ = self.lstm(x)
        x = self.fc(x[:, -1, :])  # 取最后一个时间步的输出
        return x

class AudioTextModel(nn.Module):
    def __init__(self, audio_model, text_model):
        super(AudioTextModel, self).__init__()
        self.audio_model = audio_model
        self.text_model = text_model
        self.fc = nn.Linear(128 * 2, 1)  # 假设输出为匹配程度的概率

    def forward(self, audio, text):
        audio_out = self.audio_model(audio)
        text_out = self.text_model(text)
        combined = torch.cat((audio_out, text_out), dim=1)
        output = self.fc(combined)
        return output

# 实例化模型
audio_model = AudioCNN()
text_model = TextEmbedding(embedding_dim=100, vocab_size=10000, text_length=100)  # 假设参数
model = AudioTextModel(audio_model, text_model)

